import cv2
import numpy as np
import matplotlib.pyplot as plt
import os
from PIL import Image
import pytesseract
from typing import List, Tuple, Optional


class LicensePlateProcessor:
    def __init__(self):
        # 车牌地区编码字典
        self.province_dict = {
            '京': 'Beijing', '津': 'Tianjin', '沪': 'Shanghai', '渝': 'Chongqing',
            '冀': 'Hebei', '豫': 'Henan', '云': 'Yunnan', '辽': 'Liaoning',
            '黑': 'Heilongjiang', '湘': 'Hunan', '皖': 'Anhui', '鲁': 'Shandong',
            '新': 'Xinjiang', '苏': 'Jiangsu', '浙': 'Zhejiang', '赣': 'Jiangxi',
            '鄂': 'Hubei', '桂': 'Guangxi', '甘': 'Gansu', '晋': 'Shanxi',
            '蒙': 'Inner Mongolia', '陕': 'Shaanxi', '吉': 'Jilin', '闽': 'Fujian',
            '贵': 'Guizhou', '粤': 'Guangdong', '青': 'Qinghai', '藏': 'Tibet',
            '川': 'Sichuan', '宁': 'Ningxia', '琼': 'Hainan'
        }

        # 英文字母和数字字典
        self.char_dict = {
            'A': 'A', 'B': 'B', 'C': 'C', 'D': 'D', 'E': 'E', 'F': 'F', 'G': 'G',
            'H': 'H', 'J': 'J', 'K': 'K', 'L': 'L', 'M': 'M', 'N': 'N', 'P': 'P',
            'Q': 'Q', 'R': 'R', 'S': 'S', 'T': 'T', 'U': 'U', 'V': 'V', 'W': 'W',
            'X': 'X', 'Y': 'Y', 'Z': 'Z',
            '0': '0', '1': '1', '2': '2', '3': '3', '4': '4',
            '5': '5', '6': '6', '7': '7', '8': '8', '9': '9'
        }

        # Tesseract OCR配置
        self.tesseract_config = '--psm 10 --oem 3 -c tessedit_char_whitelist=ABCDEFGHJKLMNPQRSTUVWXYZ0123456789'

    def load_image(self, image_path: str) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """加载并预处理图像"""
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"无法读取图像: {image_path}")

        # 转换为RGB用于显示和HSV用于颜色过滤
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

        return img, img_rgb, img_hsv

    def preprocess_image(self, img: np.ndarray, img_hsv: np.ndarray) -> np.ndarray:
        """多通道图像预处理，结合颜色过滤、边缘检测和形态学操作"""
        # 1. 颜色过滤 - 检测蓝色、绿色和黄色车牌
        lower_blue = np.array([100, 100, 100])
        upper_blue = np.array([130, 255, 255])
        blue_mask = cv2.inRange(img_hsv, lower_blue, upper_blue)

        lower_green = np.array([35, 100, 100])
        upper_green = np.array([90, 255, 255])
        green_mask = cv2.inRange(img_hsv, lower_green, upper_green)

        lower_yellow = np.array([15, 100, 100])
        upper_yellow = np.array([35, 255, 255])
        yellow_mask = cv2.inRange(img_hsv, lower_yellow, upper_yellow)

        # 合并所有颜色掩码
        color_mask = cv2.bitwise_or(blue_mask, green_mask)
        color_mask = cv2.bitwise_or(color_mask, yellow_mask)

        # 2. 边缘检测
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        # 高斯模糊
        blur = cv2.GaussianBlur(gray, (5, 5), 0)

        # 使用Sobel算子计算梯度
        sobelx = cv2.Sobel(blur, cv2.CV_64F, 1, 0, ksize=3)
        sobely = cv2.Sobel(blur, cv2.CV_64F, 0, 1, ksize=3)
        gradient = cv2.magnitude(sobelx, sobely)
        gradient = cv2.normalize(gradient, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)

        # 阈值处理
        _, edges = cv2.threshold(gradient, 50, 255, cv2.THRESH_BINARY)

        # 3. 合并颜色和边缘信息
        combined = cv2.bitwise_or(edges, color_mask)

        # 4. 形态学操作增强区域
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
        processed = cv2.morphologyEx(combined, cv2.MORPH_CLOSE, kernel, iterations=2)
        processed = cv2.morphologyEx(processed, cv2.MORPH_OPEN, kernel, iterations=1)

        return processed

    def locate_plate(self, img: np.ndarray, processed: np.ndarray) -> List[np.ndarray]:
        """多尺度车牌定位，提高检测成功率"""
        # 查找轮廓
        contours, _ = cv2.findContours(processed.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        # 根据轮廓面积和形状过滤可能的车牌区域
        possible_plates = []
        img_height, img_width = img.shape[:2]

        for contour in contours:
            # 计算轮廓周长
            perimeter = cv2.arcLength(contour, True)
            # 用多边形近似轮廓
            approx = cv2.approxPolyDP(contour, 0.02 * perimeter, True)

            # 车牌通常是四边形或近似四边形
            if 4 <= len(approx) <= 10:
                x, y, w, h = cv2.boundingRect(approx)

                # 过滤具有车牌长宽比的区域
                aspect_ratio = w / h
                area = w * h

                # 多级过滤条件，提高灵活性
                if (2.0 <= aspect_ratio <= 7.0 and
                        0.003 * img_width * img_height <= area <= 0.15 * img_width * img_height):

                    # 计算轮廓紧致度
                    compactness = 4 * np.pi * cv2.contourArea(contour) / (perimeter * perimeter)

                    # 车牌通常具有中等紧致度
                    if 0.05 <= compactness <= 0.9:
                        possible_plates.append(approx)

        # 如果没有找到候选，尝试降低标准
        if not possible_plates:
            for contour in contours:
                perimeter = cv2.arcLength(contour, True)
                approx = cv2.approxPolyDP(contour, 0.02 * perimeter, True)

                if 4 <= len(approx) <= 10:
                    x, y, w, h = cv2.boundingRect(approx)
                    aspect_ratio = w / h
                    area = w * h

                    if (1.5 <= aspect_ratio <= 8.0 and
                            0.002 * img_width * img_height <= area <= 0.2 * img_width * img_height):
                        possible_plates.append(approx)

        return possible_plates

    def correct_plate_perspective(self, img: np.ndarray, plate_contour: np.ndarray) -> np.ndarray:
        """矫正车牌透视，处理各种轮廓形状"""
        # 获取轮廓的四个顶点
        if len(plate_contour) == 4:
            # 如果已经是四边形，直接使用
            pts = plate_contour.reshape(4, 2)
        else:
            # 对于非四边形轮廓，计算最小面积矩形
            rect = cv2.minAreaRect(plate_contour)
            pts = cv2.boxPoints(rect)

        # 确定顶点顺序
        rect = np.zeros((4, 2), dtype="float32")

        # 计算四个顶点的和
        s = pts.sum(axis=1)
        # 左上角点的和最小，右下角点的和最大
        rect[0] = pts[np.argmin(s)]
        rect[2] = pts[np.argmax(s)]

        # 计算四个顶点的差
        diff = np.diff(pts, axis=1)
        # 右上角点的差最小，左下角点的差最大
        rect[1] = pts[np.argmin(diff)]
        rect[3] = pts[np.argmax(diff)]

        # 计算目标矩形的宽度和高度
        (tl, tr, br, bl) = rect
        widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
        widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
        maxWidth = max(int(widthA), int(widthB))

        heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
        heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
        maxHeight = max(int(heightA), int(heightB))

        # 定义目标矩形的四个顶点，调整长宽比以匹配真实车牌
        aspect_ratio = maxWidth / maxHeight
        if aspect_ratio < 3.0:
            # 如果长宽比太小，调整宽度
            maxWidth = int(maxHeight * 3.5)

        dst = np.array([
            [0, 0],
            [maxWidth - 1, 0],
            [maxWidth - 1, maxHeight - 1],
            [0, maxHeight - 1]], dtype="float32")

        # 计算透视变换矩阵并应用变换
        M = cv2.getPerspectiveTransform(rect, dst)
        warped = cv2.warpPerspective(img, M, (maxWidth, maxHeight))

        return warped

    def remove_plate_border(self, plate_img: np.ndarray) -> np.ndarray:
        """去除车牌边框，增强字符提取效果"""
        # 转换为灰度图
        gray = cv2.cvtColor(plate_img, cv2.COLOR_BGR2GRAY)

        # 自适应阈值处理
        thresh = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)

        # 形态学操作去除小噪声
        kernel = np.ones((2, 2), np.uint8)
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)

        # 查找字符区域
        contours, _ = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # 创建掩码
        mask = np.zeros(thresh.shape, dtype=np.uint8)

        # 过滤有效轮廓并在掩码上绘制
        for cnt in contours:
            x, y, w, h = cv2.boundingRect(cnt)
            # 过滤合理的字符区域
            if (thresh.shape[0] * 0.4 < h < thresh.shape[0] * 0.95 and
                    thresh.shape[1] * 0.02 < w < thresh.shape[1] * 0.3):
                cv2.rectangle(mask, (x, y), (x + w, y + h), (255, 255, 255), -1)

        # 膨胀掩码以连接可能的字符部分
        kernel = np.ones((3, 3), np.uint8)
        mask = cv2.dilate(mask, kernel, iterations=1)

        # 将掩码应用于原始图像
        result = cv2.bitwise_and(plate_img, plate_img, mask=mask)

        return result

    def segment_characters(self, plate_img: np.ndarray, output_dir: Optional[str] = None,
                           visualize: bool = False) -> Tuple[List[np.ndarray], List[Tuple[int, int, int, int]]]:
        """
        改进的字符分割方法，处理各种光照和模糊条件

        Args:
            plate_img: 车牌图像
            output_dir: 字符保存目录，如果为None则不保存
            visualize: 是否可视化字符分割结果

        Returns:
            字符图像列表和对应的边界框列表
        """
        # 转换为灰度图
        gray = cv2.cvtColor(plate_img, cv2.COLOR_BGR2GRAY)

        # 自适应阈值处理
        thresh = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)

        # 去除小噪声
        kernel = np.ones((2, 2), np.uint8)
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)

        # 水平投影确定字符的上下边界
        horizontal_projection = np.sum(thresh, axis=1) / 255

        # 确定字符边界
        height = thresh.shape[0]
        upper_bound = 0
        lower_bound = height - 1

        # 找到上边界
        for i in range(height):
            if horizontal_projection[i] > height * 0.1:
                upper_bound = i
                break

        # 找到下边界
        for i in range(height - 1, -1, -1):
            if horizontal_projection[i] > height * 0.1:
                lower_bound = i
                break

        # 裁剪上下边界
        if upper_bound < lower_bound:
            thresh_cropped = thresh[upper_bound:lower_bound + 1, :]
        else:
            thresh_cropped = thresh.copy()

        # 垂直投影确定字符的左右边界
        vertical_projection = np.sum(thresh_cropped, axis=0) / 255

        # 确定字符边界
        character_bounds = []
        in_character = False
        start = 0

        for i in range(len(vertical_projection)):
            if vertical_projection[i] > 5 and not in_character:
                in_character = True
                start = i
            elif vertical_projection[i] <= 5 and in_character:
                in_character = False
                # 过滤非常小的区域
                if i - start > thresh_cropped.shape[1] * 0.03:
                    character_bounds.append((start, i))

        # 处理连接的字符
        processed_bounds = []
        for (start, end) in character_bounds:
            char_width = end - start
            # 如果字符宽度超过平均宽度的1.8倍，可能是连接的字符
            if char_width > thresh_cropped.shape[1] * 0.18:
                # 尝试分割连接的字符
                mid = (start + end) // 2
                processed_bounds.append((start, mid))
                processed_bounds.append((mid, end))
            else:
                processed_bounds.append((start, end))

        # 提取字符
        characters = []
        char_boxes = []  # 存储字符的边界框信息 (x, y, w, h)
        char_index = 1  # 用于记录字符序号

        # 创建可视化图像
        if visualize:
            vis_img = plate_img.copy()
            vis_img = cv2.cvtColor(vis_img, cv2.COLOR_BGR2RGB)

        for (start, end) in processed_bounds:
            x = start
            y = upper_bound
            w = end - start
            h = lower_bound - upper_bound + 1

            char_img = thresh_cropped[:, start:end]
            # 确保字符图像有合理的高度和宽度
            if char_img.shape[0] > 10 and char_img.shape[1] > 5:
                # 调整字符图像大小
                char_img = cv2.resize(char_img, (50, 100))
                characters.append(char_img)
                char_boxes.append((x, y, w, h))

                # 保存每个字符图像
                if output_dir:
                    os.makedirs(output_dir, exist_ok=True)
                    # 添加字符识别结果到文件名
                    try:
                        pil_img = Image.fromarray(char_img)
                        if char_index == 1:  # 第一个字符是省份
                            char_ocr = pytesseract.image_to_string(pil_img, lang='chi_sim', config='--psm 10').strip()
                        else:  # 其他字符是字母或数字
                            char_ocr = pytesseract.image_to_string(pil_img, config=self.tesseract_config).strip()
                        # 确保文件名合法
                        valid_char = char_ocr if char_ocr.isalnum() else f"char_{char_index}"
                        filename = f"char_{valid_char}_{char_index}.png"
                    except:
                        filename = f"char_{char_index}.png"

                    cv2.imwrite(os.path.join(output_dir, filename), char_img)

                # 在可视化图像上绘制边界框
                if visualize:
                    cv2.rectangle(vis_img, (x, y), (x + w, y + h), (0, 255, 0), 2)
                    cv2.putText(vis_img, f"{char_index}", (x, y - 5),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

                char_index += 1

        # 显示可视化结果
        if visualize and len(characters) > 0:
            plt.figure(figsize=(10, 6))
            plt.title('字符分割结果')
            plt.imshow(vis_img)
            plt.axis('off')
            plt.show()

        return characters, char_boxes

    def recognize_plate_text(self, characters: List[np.ndarray]) -> str:
        """使用OCR识别车牌字符"""
        if not characters:
            return "No characters to recognize"

        plate_text = ""

        # 第一个字符通常是省份简称
        province_char = characters[0]
        pil_img = Image.fromarray(province_char)
        try:
            # 尝试使用OCR识别省份字符
            province_ocr = pytesseract.image_to_string(pil_img, lang='chi_sim', config='--psm 10').strip()
            # 如果OCR识别结果在字典中，使用字典映射；否则保留原始识别结果
            if province_ocr in self.province_dict:
                plate_text += f"{province_ocr}({self.province_dict[province_ocr]}) "
            else:
                plate_text += f"{province_ocr}? "
        except Exception as e:
            plate_text += "? "
            print(f"Error recognizing province character: {e}")

        # 剩余字符是字母和数字
        for char_img in characters[1:]:
            pil_img = Image.fromarray(char_img)
            try:
                char_ocr = pytesseract.image_to_string(pil_img, config=self.tesseract_config).strip()
                if char_ocr and char_ocr in self.char_dict:
                    plate_text += self.char_dict[char_ocr]
                else:
                    plate_text += "?"
            except Exception as e:
                plate_text += "?"
                print(f"Error recognizing character: {e}")

        return plate_text.strip()

    def visualize_results(self, original_img: np.ndarray, plate_img: np.ndarray, characters: List[np.ndarray],
                          plate_text: str, output_dir: Optional[str] = None,
                          char_boxes: Optional[List[Tuple[int, int, int, int]]] = None):
        """可视化处理结果并保存字符图像"""
        plt.figure(figsize=(15, 10))

        # 计算需要的子图总数
        total_subplots = 2 + max(1, len(characters))  # 至少3个子图

        # 显示原始图像
        plt.subplot(2, 3, 1)
        plt.title('Original Image')
        plt.imshow(cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB))
        plt.axis('off')

        # 显示定位的车牌和识别的文本
        plt.subplot(2, 3, 2)
        plt.title(f'License Plate - Recognized: {plate_text}')
        plt.imshow(cv2.cvtColor(plate_img, cv2.COLOR_BGR2RGB))
        plt.axis('off')

        # 显示分割的字符
        if characters:
            for i, char in enumerate(characters):
                plt.subplot(2, 3, i + 3)
                plt.title(f'Character {i + 1}')
                plt.imshow(char, cmap='gray')
                plt.axis('off')

                # 保存字符图像
                if output_dir and char_boxes:
                    x, y, w, h = char_boxes[i]
                    # 从原始车牌图像中提取字符（彩色版本）
                    char_color = plate_img[y:y + h, x:x + w]
                    if char_color.size > 0:
                        char_color = cv2.resize(char_color, (50, 100))
                        cv2.imwrite(os.path.join(output_dir, f'char_color_{i + 1}.png'), char_color)
        else:
            # 如果没有找到字符，显示消息
            plt.subplot(2, 3, 3)
            plt.title('No Characters Segmented')
            plt.text(0.5, 0.5, 'Failed to segment characters', ha='center', va='center')
            plt.axis('off')

        plt.tight_layout()

        # 保存可视化结果
        if output_dir:
            plt.savefig(os.path.join(output_dir, 'results_visualization.png'))

        plt.show()

    def process_license_plate(self, image_path: str, output_dir: Optional[str] = None,
                              visualize: bool = True) -> Tuple[str, np.ndarray, List[np.ndarray]]:
        """
        处理车牌图像的完整流程

        Args:
            image_path: 输入图像路径
            output_dir: 输出目录，如果为None则不保存
            visualize: 是否可视化结果

        Returns:
            识别的车牌文本，车牌图像，字符图像列表
        """
        try:
            print(f"Processing image: {image_path}")

            # 加载图像
            img, img_rgb, img_hsv = self.load_image(image_path)

            # 预处理图像
            processed = self.preprocess_image(img, img_hsv)

            # 定位车牌
            plate_contours = self.locate_plate(img, processed)

            if plate_contours:
                print(f"Found {len(plate_contours)} possible license plate regions")

                # 选择最大的轮廓作为车牌
                plate_contour = max(plate_contours, key=lambda c: cv2.contourArea(c))

                # 在原图上绘制车牌轮廓
                img_with_contour = img.copy()
                cv2.drawContours(img_with_contour, [plate_contour], -1, (0, 255, 0), 2)

                # 矫正车牌透视
                plate_img = self.correct_plate_perspective(img, plate_contour)

                # 去除车牌边框
                plate_img_without_border = self.remove_plate_border(plate_img)

                # 分割字符
                characters, char_boxes = self.segment_characters(
                    plate_img_without_border, output_dir, visualize)

                print(f"Successfully segmented {len(characters)} characters")

                # 识别车牌文本
                plate_text = self.recognize_plate_text(characters)
                print(f"Recognized license plate: {plate_text}")

                # 可视化结果并保存字符图像
                if visualize:
                    self.visualize_results(img_with_contour, plate_img_without_border,
                                           characters, plate_text, output_dir, char_boxes)

                # 保存处理后的车牌图像
                if output_dir:
                    cv2.imwrite(os.path.join(output_dir, 'plate_result.jpg'), plate_img_without_border)
                    print(f"Results saved to {output_dir}")

                return plate_text, plate_img_without_border, characters
            else:
                print("No license plate found. Try adjusting parameters or using a different image.")

                # 显示预处理后的图像和原图，帮助调试
                if visualize:
                    plt.figure(figsize=(15, 10))

                    plt.subplot(121)
                    plt.title('Original Image')
                    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
                    plt.axis('off')

                    plt.subplot(122)
                    plt.title('Processed Image')
                    plt.imshow(processed, cmap='gray')
                    plt.axis('off')

                    plt.tight_layout()
                    plt.show()

                return "No plate detected", None, []

        except Exception as e:
            print(f"Error during processing: {e}")
            return f"Error: {str(e)}", None, []


# 使用示例
if __name__ == "__main__":
    # 创建车牌处理器实例
    processor = LicensePlateProcessor()

    # 要处理的图像路径
    image_path = '3.png'  # 替换为实际的图像路径
    output_dir = 'character_images'  # 字符保存目录

    try:
        # 处理车牌图像
        plate_text, plate_img, characters = processor.process_license_plate(
            image_path, output_dir, visualize=True)

        # 打印结果
        if plate_text and plate_img is not None:
            print("\n=== License Plate Recognition Results ===")
            print(f"Recognized Plate: {plate_text}")
            print(f"Number of characters segmented: {len(characters)}")
            if output_dir:
                print(f"Individual characters saved to: {output_dir}")

    except Exception as e:
        print(f"Failed to process image: {e}")